r"""undocumented
轻量封装的 Pytorch LSTM 模块.
可在 forward 时传入序列的长度, 自动对padding做合适的处理.
"""

__all__ = [
    "LSTM"
]

import torch
import torch.nn as nn
import torch.nn.utils.rnn as rnn
# from .mask_lstm import mask_lstm
from . import lstm


class LSTM(nn.Module):
    def __init__(self, input_size, hidden_size=100, num_layers=1, dropout=0.0, batch_first=True,
                 bidirectional=False, bias=True, nlp_mask=True, threshold_pruning_method=True):
        r"""

        :param input_size:  输入 `x` 的特征维度
        :param hidden_size: 隐状态 `h` 的特征维度. 如果bidirectional为True，则输出的维度会是hidde_size*2
        :param num_layers: rnn的层数. Default: 1
        :param dropout: 层间dropout概率. Default: 0
        :param bidirectional: 若为 ``True``, 使用双向的RNN. Default: ``False``
        :param batch_first: 若为 ``True``, 输入和输出 ``Tensor`` 形状为
            :(batch, seq, feature). Default: ``False``
        :param bias: 如果为 ``False``, 模型将不会使用bias. Default: ``True``
        """
        super(LSTM, self).__init__()
        self.batch_first = batch_first
        self.lstm = lstm.LSTM(input_size, hidden_size, num_layers, bias=bias, batch_first=batch_first,
                            dropout=dropout, bidirectional=bidirectional, nlp_mask=nlp_mask, threshold_pruning_method=threshold_pruning_method)
        self.init_param()

    def init_param(self):
        for name, param in self.named_parameters():
            if 'bias' in name:
                param.data.fill_(0)
                n = param.size(0)
                start, end = n // 4, n // 2
                param.data[start:end].fill_(1)
            elif 'score' in name:
                pass
            else:
                nn.init.xavier_uniform_(param)

    def forward(self, x, seq_len=None, h0=None, c0=None, task_id=None, sparsity=0.5, epoch=None):
        # for name, param in self.named_parameters():
        #     print(name)

        batch_size, max_len, _ = x.size()
        if h0 is not None and c0 is not None:
            hx = (h0, c0)
        else:
            hx = None
        if seq_len is not None and not isinstance(x, rnn.PackedSequence):
            sort_lens, sort_idx = torch.sort(seq_len, dim=0, descending=True)
            if self.batch_first:
                x = x[sort_idx]
            else:
                x = x[:, sort_idx]
            x = rnn.pack_padded_sequence(x, sort_lens, batch_first=self.batch_first)
            output, hx = self.lstm(x, hx, task_id = task_id, sparsity=sparsity, epoch=epoch)  # -> [N,L,C]
            output, _ = rnn.pad_packed_sequence(output, batch_first=self.batch_first, total_length=max_len)
            _, unsort_idx = torch.sort(sort_idx, dim=0, descending=False)
            if self.batch_first:
                output = output[unsort_idx]
            else:
                output = output[:, unsort_idx]
            hx = hx[0][:, unsort_idx], hx[1][:, unsort_idx]
        else:
            output, hx = self.lstm(x, hx, task_id = task_id)
        return output, hx
